import os
import pdfplumber
import re
from pathlib import Path
from datetime import datetime
import json
from openpyxl import Workbook

def extract_pdf_text(pdf_path):
    """使用pdfplumber提取PDF内容"""
    try:
        header_text = []
        table_text = []
        with pdfplumber.open(pdf_path) as pdf:
            for i, page in enumerate(pdf.pages, 1):
                try:
                    # 提取表格
                    tables = page.extract_tables()
                    table_bounds = []
                    if tables:
                        # 记录所有表格的位置
                        for table in tables:
                            if table:
                                table_bounds.append(page.find_tables()[0].bbox)
                                # 处理表格内容
                                for row in table:
                                    if not any(cell for cell in row):
                                        continue
                                    line = '  '.join(cell.strip() if cell else '' for cell in row)
                                    if line.strip():
                                        table_text.append(line)
                    
                    # 提取表格上方的文字
                    if table_bounds:
                        top_bound = min(bound[1] for bound in table_bounds)
                        words = page.extract_words()
                        header_words = [w for w in words if w['top'] < top_bound]
                        
                        # 按行组织文字
                        current_line = []
                        current_top = None
                        for word in sorted(header_words, key=lambda x: (x['top'], x['x0'])):
                            if current_top is None:
                                current_top = word['top']
                            elif abs(word['top'] - current_top) > 5:  # 新行
                                header_text.append(' '.join(current_line))
                                current_line = []
                                current_top = word['top']
                            current_line.append(word['text'])
                        
                        if current_line:
                            header_text.append(' '.join(current_line))
                
                except Exception as e:
                    print(f"第 {i} 页处理失败：{str(e)}")
        
        # 合并文本，确保表格上方文字在前
        all_text = header_text + table_text
        return '\n'.join(all_text) if all_text else None

    except Exception as e:
        print(f"文件打开失败：{str(e)}")
        return None

def save_to_excel(items, output_path):
    """将发票数据保存为Excel文件"""
    wb = Workbook()
    ws = wb.active
    ws.title = "发票数据"
    
    # 设置表头
    headers = [
        "发票类型", "发票号码", "开票日期", "购买方名称", 
        "销售方名称", "销售方税号", "商品名称", "金额",
        "税率", "税额", "单位", "备注"
    ]
    ws.append(headers)
    
    # 写入数据
    for item in items:
        row = [
            item['type'],
            item['number'],
            item['date'],
            item['buyer_name'],
            item['seller_name'],
            item['seller_tax_id'],
            item['name'],
            item['amount'],
            item['tax_rate'],
            item['tax'],
            item['unit'] or '',  # 处理 None 值
            item['notes']
        ]
        ws.append(row)
    
    # 调整列宽
    for col in ws.columns:
        max_length = 0
        column = col[0].column_letter
        for cell in col:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = (max_length + 2)
        ws.column_dimensions[column].width = adjusted_width
    
    # 保存文件
    wb.save(output_path)

def parse_invoice_text(text):
    """解析发票文本生成结构化数据"""
    # 初始化结构
    result = {
        "invoice_base": {},
        "buyer": {},
        "seller": {},
        "items": [],
        "total": {},
        "notes": ""
    }
    
    lines = [line.strip() for line in text.split('\n') if line.strip()]
    
    # 解析基础信息
    for i, line in enumerate(lines):
        if '电子发票' in line:
            result['invoice_base']['type'] = line.replace('(', '（').split('（')[1].split('）')[0]
        elif '发票号码' in line:
            result['invoice_base']['number'] = line.split('：')[1].strip()
        elif '开票日期' in line:
            try:
                date_str = line.split('：')[1].strip()
                date_str = (date_str.replace('年', '-')
                                  .replace('月', '-')
                                  .replace('日', ''))
                result['invoice_base']['date'] = datetime.strptime(date_str, "%Y-%m-%d").isoformat()
            except Exception as e:
                print(f"日期解析失败：{str(e)}")
        elif '开票人' in line:
            result['invoice_base']['issuer'] = line.split('：')[-1].strip()

    # 解析购买方和销售方信息
    for i, line in enumerate(lines):
        if '名称：' in line:
            if '1#' in line:  # 购买方
                result['buyer']['name'] = line.split('名称：')[1].strip()
            elif '泰禾' in line:  # 销售方
                result['seller']['name'] = line.split('名称：')[1].strip()
        
        # 解析税号
        if '统一社会信用代码' in line:
            next_line = lines[i + 1] if i + 1 < len(lines) else ''
            if '统一社会信用代码' in next_line:
                result['seller']['tax_id'] = line.split('：')[1].strip()
                result['seller']['tax_id'] = next_line.split('：')[1].strip()
            else:
                result['seller']['tax_id'] = line.split('：')[1].strip()

    # 解析商品明细
    items = []
    current_item = None
    
    for i, line in enumerate(lines):
        if line.startswith('*'):
            # 如果有未完成的项目，先保存
            if current_item and 'amount' in current_item:
                items.append(current_item)
            
            # 提取商品名称
            name_parts = line.strip('*').split('*')
            current_item = {'name': name_parts[-1].strip().split(' ')[0]}
            
            # 检查同一行是否包含金额信息
            amount_parts = re.findall(r'([\d.]+)\s+(\d+)%\s+([\d.]+)', line)
            if amount_parts:
                amount, rate, tax = amount_parts[0]
                current_item.update({
                    'amount': float(amount),
                    'tax_rate': rate,
                    'tax': float(tax),
                    'unit': '月' if '月' in line else None
                })
                items.append(current_item)
                current_item = None
        elif current_item and not 'amount' in current_item:
            # 处理跨行的金额信息
            amount_parts = re.findall(r'([\d.]+)\s+(\d+)%\s+([\d.]+)', line)
            if amount_parts:
                amount, rate, tax = amount_parts[0]
                current_item.update({
                    'amount': float(amount),
                    'tax_rate': rate,
                    'tax': float(tax),
                    'unit': '月' if '月' in line else None
                })
                items.append(current_item)
                current_item = None
            elif '月' in line and not any(c.isdigit() for c in line):
                current_item['unit'] = '月'
    
    # 确保最后一个项目被添加
    if current_item and 'amount' in current_item:
        items.append(current_item)

    # 过滤并去重商品明细
    seen_items = set()
    result['items'] = []
    for item in items:
        item_key = f"{item['name']}_{item['amount']}"
        if item_key not in seen_items:
            result['items'].append(item)
            seen_items.add(item_key)

    # 解析总计金额
    for line in lines:
        if '合 计' in line:
            amounts = re.findall(r'¥([\d.]+)', line)
            if len(amounts) >= 2:
                result['total'].update({
                    'amount_without_tax': float(amounts[0]),
                    'tax_amount': float(amounts[1])
                })
        elif '价税合计' in line:
            amount = re.findall(r'¥([\d.]+)', line)
            if amount:
                result['total']['total_amount'] = float(amount[0])

    # 解析备注
    for i, line in enumerate(lines):
        if '备' in line and i + 1 < len(lines) and '注' in lines[i + 1]:
            # 处理第一行备注（截取"注"后面的内容）
            notes = [lines[i + 1].split('注')[-1].strip()]
            
            # 处理后续行的备注
            j = i + 2
            while j < len(lines):
                next_line = lines[j].strip()
                if not next_line or '开票人' in next_line:
                    break
                notes.append(next_line)
                j += 1
            
            # 合并所有备注为一个字符串
            result['notes'] = ' '.join(notes).strip()
            break

    return result

# 在batch_process函数中调用
def batch_process(input_dir):
    """批量处理目录下的PDF文件"""
    input_path = Path(input_dir)
    output_dir = input_path / 'pdfplumber_texts'
    output_dir.mkdir(exist_ok=True)
    
    pdf_files = list(input_path.glob('**/*.pdf'))
    if not pdf_files:
        print("未找到PDF文件")
        return
    
    print(f"发现 {len(pdf_files)} 个PDF文件")
    # 将 items 转换为一维数组格式
    flattened_items = []
    for i, pdf_file in enumerate(pdf_files, 1):
        print(f"\n处理进度 ({i}/{len(pdf_files)})")
        text = extract_pdf_text(pdf_file)
        
        if text:
            output_file = output_dir / f"{pdf_file.stem}.txt"
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(text)
            print(f"已保存：{output_file}")
    
            result = parse_invoice_text(text)
            
            
            for item in result['items']:
                flat_item = {
                    "type": result['invoice_base']['type'],
                    "number": result['invoice_base']['number'],
                    "date": result['invoice_base']['date'],
                    "buyer_name": result['buyer']['name'],
                    "seller_name": result['seller']['name'],
                    "seller_tax_id": result['seller']['tax_id'],
                    "notes": result['notes'],
                    'name': item['name'],
                    'amount': item['amount'],
                    'tax_rate': item['tax_rate'],
                    'tax': item['tax'],
                    'unit': item['unit']
                }
                flattened_items.append(flat_item)
            
            json_data = json.dumps(flattened_items, ensure_ascii=False, indent=2)
            output_file = output_dir / f"{pdf_file.stem}.json"
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(json_data)
                print(f"已保存：{output_file}")
    
    # 在保存完 JSON 文件后添加
    excel_output = output_dir / f"{pdf_file.stem}.xlsx"
    save_to_excel(flattened_items, excel_output)
    print(f"已保存：{excel_output}")

if __name__ == "__main__":
    input_dir = input("请输入PDF目录路径（留空使用当前目录）：").strip() or os.getcwd()
    batch_process(input_dir)
    print("\n处理完成！")